
"""
获取数据表
"""
from tkinter.constants import FALSE

from core.state import AnalyzeState
from langgraph.types import Command
from langgraph.config import get_stream_writer
from core.utils.mcp_function import call_tool
from core.llm import dataset_usage_analyzer
from models.aiquery import AiQueryData, TableData
import concurrent.futures
from config import settings
from loguru import logger
import pandas as pd
import asyncio


async def biw_query_node(state: AnalyzeState):

    writer = get_stream_writer()
    writer("智能问数查询接口获取数据")
    logger.info("智能问数查询接口获取数据")

    session_id = state.get('sessionId')
    queries = state.get('rewrites')
    biw_tasks = [biw_query(session_id, query) for query in queries]

    # 并发执行智能问数查询
    # datas = await asyncio.gather(*biw_tasks)
    datas = []
    for task in asyncio.as_completed(biw_tasks):
        try:
            biw_result = await task
            query = biw_result.get('query')
            data = biw_result.get('data')
            if not isinstance(data, dict):
                logger.error(f'查询失败：{query}. (返回结果：{data})')
            elif 'aiResult' not in data or 'dataQueryReq' not in data:
                logger.error(f'查询失败：{query}. (结果不符合格式要求：{data})')
            else:
                logger.info(f'查询成功：{query}.')
                datas.append((query, data))
        except Exception as e:
            logger.error(e)

    # 智能问数查询的数据可能含多个数据集，故线程池处理。
    datasets = state.get('datasets')
    with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
        tasks = []
        for (query, data) in datas:
            columns = data['dataQueryReq']['indicatorTable']['columns']
            for dataset in data['aiResult']['datasets']:
                rows = dataset.pop('rows', [])
                if not isinstance(rows, list) or len(rows) == 0:
                    continue
                if len(rows) == 1 and len(rows[0]) == 0:
                    continue

                schema = dict()
                if 'table' in dataset and isinstance(dataset['table'], dict):
                    schema['name'] = dataset['table'].get('name', '')
                else:
                    schema['name'] = dataset.get('name', '')
                schema['query'] = query
                schema['sql'] = dataset.get('sql', '')

                # 获取数据表的字段信息fields。
                df = pd.DataFrame(rows[0:100])
                keys = df.columns.tolist()
                if len(keys) == 0:
                    keys = dataset['dimensions']
                if 'table' in dataset and 'columns' in dataset['table']:
                    table_columns = dataset['table']['columns']
                    schema['columns'] = [item for item in table_columns if item['name'] in keys]
                else:
                    schema['columns'] = [item for item in columns if item['name'] in keys]
                tasks.append(executor.submit(dataset_usage_analyzer.invoke,  schema, rows))

        for future in concurrent.futures.as_completed(tasks):
            if future.exception() is not None:
                continue
            result = future.result()
            if len(result) != 0:
                description = result.pop('description', '')
                datasets.append(AiQueryData(
                    description=description,
                    data=TableData(**result))
                )

    writer(datasets)
    query_datas = [data for data in datasets if isinstance(data, AiQueryData)]
    logger.info(f'智能问数获取的数据集个数：{len(query_datas)}')

    # return Command(update={"datasets": datasets}, goto='dataset_filter')

    # 增加未获取到数据的查询列表.
    used_queries = set()
    for dataset in datasets:
        if isinstance(dataset, AiQueryData) and isinstance(dataset.data, TableData):
            used_queries.add(dataset.data.query)

    rewrites = state.get('rewrites')
    for query in rewrites:
        if query not in used_queries:
            datasets.append(AiQueryData(
                description=query,
                data=None
            ))
    writer(datasets)
    logger.info(f'最终返回的数据集：{datasets}')

    return Command(update={"datasets": datasets})

async def biw_query(session_id:str, query: str):
    data = await call_tool(
        sse_url=settings.DATA_SERVER_URI,
        tool_name='system#biwQueryBySessionId',
        tool_args={
            'sessionId': session_id,
            'query': query
        }
    )
    return {
        'query': query,
        'data': data
    }
